Skip to content

Enable offloading of validation models to CPU during training runs for EMTTS - #16035

Open
blisc wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
blisc:jasoli/easymagpie-validation-offload
Open

Enable offloading of validation models to CPU during training runs for EMTTS#16035
blisc wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
blisc:jasoli/easymagpie-validation-offload

Conversation

@blisc

@blisc blisc commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do ?

Enable offloading of validation models to CPU during training runs for EMTTS

Collection: tts

Changelog

  • Enable offloading of validation models to CPU during training runs for EMTTS

PR Type:

  • New Feature
  • Bugfix
  • Documentation

Signed-off-by: Jason <jasoli@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 5, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions github-actions Bot added the TTS label Aug 5, 2026

@rfejgin rfejgin left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In general looks good, but I had some top level questions: Is the purpose to free up GPU memory to enable larger batch sizes in training? If so, could you say a bit on how much GPU memory the validation models use to start with?

)

def training_step(self, batch, batch_idx):
uses_uncached_audio = 'context_audio_codes' not in batch or 'audio_codes' not in batch

@Edresson Edresson Aug 6, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should extend this if statement to check whether the codec is required for multi-turn user-audio conditioning:

self.cfg.get("use_multiturn_dataset", False)
and batch["user_audio_turn_splitted"] is not None
and self.cfg.get("condition_on_user_speech", False)

Without this check, when context_audio_codes and audio_codes are included in the multi-turn data, user-audio conditioning feature extraction may run with the codec on the CPU, which would be extremely slow or raise errors.

audio_lens = batch['audio_lens']
audio_codes, audio_codes_lens = self._codec_helper.audio_to_codes(audio, audio_lens)

if uses_uncached_audio and self._should_offload_validation_models():

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this necessary? Could we keep the codec on the GPU throughout an epoch whenever any batch in that epoch requires it?

If you still want to do it, please move it to the after the following block, otherwise multiturn training data will fails.

 if (
            self.cfg.get("use_multiturn_dataset", False)
            and batch["user_audio_turn_splitted"] is not None
            and self.cfg.get("condition_on_user_speech", False)
        ):

@Edresson Edresson left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this change necessary?

Also, have you measured its impact on validation performance? Validation is already relatively slow, and I’m concerned that this PR may introduce additional overhead.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants